﻿using System;
using Algoritmia.DataStructures.Maps;
using Algoritmia.SemiRings;
using Algoritmia.StateSpaces;
using Algoritmia.Utils;

namespace Algoritmia.Schemes.DynamicProgramming
{
    public class ForwardDynamicProgrammingSolver<TState, TDecision, TResult,TWeight>
    {
        private readonly Func<IForwardStateSpace<TState, TDecision>, IMap<TState, TResult>> _createMap;
        private readonly Func<StateSpaceTopsorter<TState, TDecision>> _createStateSpaceTopsorter;

        public ForwardDynamicProgrammingSolver(Func<IForwardStateSpace<TState, TDecision>, IMap<TState, TResult>> createMap = null)
        {
            _createMap = createMap ?? (space => AlgoritmiaFactory.Maps<TState, TResult>.Default);
            _createStateSpaceTopsorter = () => new StateSpaceTopsorter<TState, TDecision>();
        }

        public TResult Solve(IForwardStateSpace<TState, TDecision> space, ISemiRing<TResult,TWeight> semiRing, Func<TState, TDecision, TWeight> wfn)
        {
            var mem = _createMap(space);
            foreach(var initial in space.InitialStates)
            {
                mem[initial] = semiRing.One;
            }

            var topsorter = _createStateSpaceTopsorter();
            var best = semiRing.Zero;
            foreach (var s in topsorter.Topsorted(space))
            {
                if (!mem.Contains(s))
                {
                    mem[s] = semiRing.Zero;
                }
                foreach(var d in space.Decisions(s))
                {
                    var s2 = space.Decide(s, d);
                    if (!mem.Contains(s2))
                    {
                        mem[s2] = semiRing.Zero;
                    }
                    mem[s2] = semiRing.Plus(mem[s2], semiRing.Times(mem[s], wfn(s, d)));
                }
                if (space.IsFinal(s))
                {
                    best = semiRing.Plus(best, mem[s]);
                }
                mem.Remove(s);
            }
            return best;
        }
    }

    public class ForwardDynamicProgrammingSolver<TState,TDecision,TResult> :
        ForwardDynamicProgrammingSolver<TState, TDecision, TResult, TResult>
    {
        public ForwardDynamicProgrammingSolver(Func<IForwardStateSpace<TState, TDecision>, IMap<TState, TResult>> createMap = null)
            : base(createMap)
        {
            
        }
    }
}

